RED = 0
BLACK = 1

class RedBlackTreeNode(object):
    def __init__(self, value):
        self.value = value
        self.left = None
        self.right = None
        self.p = None   #父亲指针
        self.color = RED

class RedBlackTree(object):
    def __init__(self):
        self.root:RedBlackTreeNode = None

    #插入二叉排序树
    def insert(self, node: RedBlackTreeNode):
        if not self.root:
            self.root = node
        else:
            x = self.root
            y = None
            while x:    #x为None时结束，代表着找到了要插入的位置，用y记录该位置的父节点
                y = x
                if node.value > x.value:
                    x = x.right
                else:
                    x = x.left

            node.p = y
            if node.value > y.value:
                y.right = node
            else:
                y.left = node

        self.insert_fixup(node)


    #颜色调整和插入后旋转
    def insert_fixup(self, node: RedBlackTreeNode):
        parent:RedBlackTreeNode = node.p
        #父亲存在且父亲颜色为红色
        while parent and parent.color == RED:
            grandparent = parent.p  #如果父亲是存在的并且父亲的颜色是红色，则爷爷节点一定存在并且是黑色
            if parent == grandparent.left:
                uncle = grandparent.right

                #判断情形三
                if uncle and uncle.color == RED:
                    parent.color = BLACK
                    uncle.color = BLACK
                    grandparent.color = RED
                    node = grandparent  #爷爷视为新插入的节点进行调整
                    parent = node.p
                    continue

                #判断情形四
                if node == parent.right:
                    left_rotate(self,node.p)    #左旋成为情形五
                    node, parent = parent, node

                #如果没有continue，则必须做情形五
                right_rotate(self,grandparent)
                parent.color = BLACK
                grandparent.color = RED
                break

            else:
                uncle = grandparent.left

                # 判断情形三
                if uncle and uncle.color == RED:
                    parent.color = BLACK
                    uncle.color = BLACK
                    grandparent.color = RED
                    node = grandparent  # 爷爷视为新插入的节点进行调整
                    parent = node.p
                    continue

                # 判断情形四
                if node == parent.left:
                    right_rotate(self, node.p)  # 左旋成为情形五
                    node, parent = parent, node

                # 如果没有continue，则必须做情形五
                left_rotate(self, grandparent)
                parent.color = BLACK
                grandparent.color = RED
                break

        self.root.color = BLACK


def left_rotate(tree:RedBlackTree, node:RedBlackTreeNode):
    if not node.right:
        return False
    node_right = node.right
    node_right.p = node.p
    if node.p is None:
        tree.root = node_right
    elif node.p.left == node:
        node.p.left = node_right
    else:
        node.p.right = node_right

    if node_right.left:
        node_right.left.p = node

    node.right = node_right.left
    node.p = node_right
    node_right.left = node

def right_rotate(tree:RedBlackTree, node:RedBlackTreeNode):
    if not node.left:
        return False
    node_left = node.left
    node_left.p = node.p
    if node.p is None:
        tree.root = node_left
    elif node.p.left == node:
        node.p.left = node_left
    else:
        node.p.right = node_left

    if node_left.right:
        node_left.right.p = node

    node.left = node_left.right
    node.p = node_left
    node_left.right = node

def rbtree_print(node:RedBlackTreeNode, key, direction):
    if node:
        if direction == 0:
            print("%2d(B) is root" % node.value)
        else:
            print("%d(%s) is %2d's %6s child" % (node.value,
                                                 ("B" if node.color == BLACK else "R"),
                                                 key,
                                                 ("right" if direction == 1 else "left")
                                                )
                  )
        rbtree_print(node.left, node.value, -1)
        rbtree_print(node.right, node.value, 1)

def main():
    number_list = [7,4,1,8,5,2,9,6,3]
    tree = RedBlackTree()
    for number in number_list:
        node = RedBlackTreeNode(number)
        tree.insert(node)
        del node
    rbtree_print(tree.root, tree.root.value, 0)
if __name__ == '__main__':
    main()